from __future__ import annotations

import logging

from chia_rs import BlockRecord, ConsensusConstants, FullBlock, HeaderBlock, UnfinishedBlock
from chia_rs.sized_bytes import bytes32
from chia_rs.sized_ints import uint8, uint32, uint64

from chia.consensus.blockchain_interface import BlockRecordsProtocol
from chia.types.unfinished_header_block import UnfinishedHeaderBlock

log = logging.getLogger(__name__)


def final_eos_is_already_included(
    header_block: UnfinishedHeaderBlock | UnfinishedBlock | HeaderBlock | FullBlock,
    blocks: BlockRecordsProtocol,
    sub_slot_iters: uint64,
) -> bool:
    """
    Args:
        header_block: An overflow block, with potentially missing information about the new sub slot
        blocks: all blocks that have been included before header_block
        sub_slot_iters: sub_slot_iters at the header_block

    Returns: True iff the missing sub slot was already included in a previous block. Returns False if the sub
    slot was not included yet, and therefore it is the responsibility of this block to include it

    """
    if len(header_block.finished_sub_slots) > 0:
        # We already have an included empty sub slot, which means the prev block is 2 sub slots behind.
        return False
    curr: BlockRecord = blocks.block_record(header_block.prev_header_hash)

    # We also check if curr is close to header_block, which means it's in the same sub slot
    seen_overflow_block = curr.overflow and (header_block.total_iters - curr.total_iters < sub_slot_iters // 2)
    while not curr.first_in_sub_slot and not curr.height == 0:
        if curr.overflow and header_block.total_iters - curr.total_iters < sub_slot_iters // 2:
            seen_overflow_block = True
        curr = blocks.block_record(curr.prev_hash)

    if curr.first_in_sub_slot and seen_overflow_block:
        # We have seen another overflow block in this slot (same as header_block), therefore there are no
        # missing sub slots
        return True

    # We have not seen any overflow blocks, therefore header_block will have to include the missing sub slot in
    # the future
    return False


def get_block_challenge(
    constants: ConsensusConstants,
    header_block: UnfinishedHeaderBlock | UnfinishedBlock | HeaderBlock | FullBlock,
    blocks: BlockRecordsProtocol,
    genesis_block: bool,
    overflow: bool,
    skip_overflow_last_ss_validation: bool,
) -> bytes32:
    if len(header_block.finished_sub_slots) > 0:
        if overflow:
            # New sub-slot with overflow block
            if skip_overflow_last_ss_validation:
                # In this case, we are missing the final sub-slot bundle (it's not finished yet), however
                # There is a whole empty slot before this block is infused
                challenge: bytes32 = header_block.finished_sub_slots[-1].challenge_chain.get_hash()
            else:
                challenge = header_block.finished_sub_slots[
                    -1
                ].challenge_chain.challenge_chain_end_of_slot_vdf.challenge
        else:
            # No overflow, new slot with a new challenge
            challenge = header_block.finished_sub_slots[-1].challenge_chain.get_hash()
    elif genesis_block:
        challenge = constants.GENESIS_CHALLENGE
    else:
        if overflow:
            if skip_overflow_last_ss_validation:
                # Overflow infusion without the new slot, so get the last challenge
                challenges_to_look_for = 1
            else:
                # Overflow infusion, so get the second to last challenge. skip_overflow_last_ss_validation is False,
                # Which means no sub slots are omitted
                challenges_to_look_for = 2
        else:
            challenges_to_look_for = 1
        reversed_challenge_hashes: list[bytes32] = []
        curr: BlockRecord = blocks.block_record(header_block.prev_header_hash)
        while len(reversed_challenge_hashes) < challenges_to_look_for:
            if curr.first_in_sub_slot:
                assert curr.finished_challenge_slot_hashes is not None
                reversed_challenge_hashes += reversed(curr.finished_challenge_slot_hashes)
                if len(reversed_challenge_hashes) >= challenges_to_look_for:
                    break
            if curr.height == 0:
                assert curr.finished_challenge_slot_hashes is not None
                assert len(curr.finished_challenge_slot_hashes) > 0
                break
            curr = blocks.block_record(curr.prev_hash)
        challenge = reversed_challenge_hashes[challenges_to_look_for - 1]
    return challenge


# Returns the previous transaction block up to the blocks signage point
# we use this for block validation since when the block is farmed we do not know the latest transaction block
# since a new one might be infused by the time the block is infused
def pre_sp_tx_block(
    constants: ConsensusConstants,
    blocks: BlockRecordsProtocol,
    *,
    prev_b_hash: bytes32,
    sp_index: uint8,
    first_in_sub_slot: bool,
) -> BlockRecord | None:
    if prev_b_hash == constants.GENESIS_CHALLENGE:
        return None
    curr = blocks.block_record(prev_b_hash)
    before_slot = first_in_sub_slot
    while curr.height > 0:
        if curr.is_transaction_block and (before_slot or curr.signage_point_index < sp_index):
            break
        if curr.first_in_sub_slot:
            before_slot = True
        curr = blocks.block_record(curr.prev_hash)
    return curr


def pre_sp_tx_block_height(
    constants: ConsensusConstants,
    blocks: BlockRecordsProtocol,
    *,
    prev_b_hash: bytes32,
    sp_index: uint8,
    first_in_sub_slot: bool,
) -> uint32:
    latest_tx_block = pre_sp_tx_block(
        constants=constants,
        blocks=blocks,
        prev_b_hash=prev_b_hash,
        sp_index=sp_index,
        first_in_sub_slot=first_in_sub_slot,
    )
    if latest_tx_block is None:
        return uint32(0)
    return latest_tx_block.height
